// Copyright 2015 Google Inc. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include <sys/socket.h>  // socket(), connect()
#include <sys/types.h>   // See socket man page
#include <sys/un.h>      // sockaddr_un
#include <errno.h>       // errno
#include <unistd.h>      // close()
#include <stdio.h>       // printf()
#include <sys/mman.h>    // mmap(), MAP_*, PROT_*
#include <stdint.h>      // uint32_t
#include <poll.h>        // poll(), POLLIN, pollfd
#include <stdlib.h>      // malloc(), free()

#include <testimony.h>

#ifdef __cplusplus
extern "C" {
#endif

#define TESTIMONY_ERRBUF_SIZE 256
#define TESTIMONY_BUF_SIZE 256

/* This list generated by go/protocol/to_c binary */
#define TESTIMONY_PROTOCOL_TYPE_BlockIndex 0
#define TESTIMONY_PROTOCOL_TYPE_ServerToClient 32769
#define TESTIMONY_PROTOCOL_TYPE_WaitingForFanoutIndex 32770
#define TESTIMONY_PROTOCOL_TYPE_FanoutSize 32771
#define TESTIMONY_PROTOCOL_TYPE_BlockSize 32772
#define TESTIMONY_PROTOCOL_TYPE_NumBlocks 32773
#define TESTIMONY_PROTOCOL_TYPE_ClientToServer 49158
#define TESTIMONY_PROTOCOL_TYPE_FanoutIndex 49159
#define TESTIMONY_PROTOCOL_TYPE_Error 65535

struct testimony_internal {
  testimony_connection conn;
  int sock_fd;
  int afpacket_fd;
  uint8_t* ring;
  char errbuf[TESTIMONY_ERRBUF_SIZE];
  uint32_t* block_counts;
  uint8_t buf[TESTIMONY_BUF_SIZE];
  uint8_t* buf_start;
  uint8_t* buf_limit;
};

static uint32_t protocol_type(uint32_t x) {
  if ((x & 0x80000000) == 0) { return TESTIMONY_PROTOCOL_TYPE_BlockIndex; }
  return x >> 16;
}

static uint32_t protocol_length(uint32_t x) {
  if ((x & 0x80000000) == 0) { return 0; }
  return x & 0xFFFF;
}

static void set_be_32(uint8_t* a, uint32_t x) {
  a[0] = x >> 24;
  a[1] = x >> 16;
  a[2] = x >> 8;
  a[3] = x;
}
#define TERR(msg, ...)                                              \
  do {                                                              \
    int err = errno;                                                \
    snprintf(t->errbuf, TESTIMONY_ERRBUF_SIZE,                      \
             msg ": %s", ##__VA_ARGS__, strerror(err));             \
    errno = err;                                                    \
  } while (0)
#define TERR_SET(errno_val, msg, ...)                               \
  do {                                                              \
    errno = errno_val;                                              \
    TERR(msg, ##__VA_ARGS__);                                       \
  } while (0)
static int recv_t(testimony t, uint8_t* out, size_t s) {
  int r;
  uint8_t* out_limit = out + s;
  size_t to_copy_out, to_copy_buf, to_copy;
  while (out < out_limit) {
    while (t->buf_start >= t->buf_limit) {
      t->buf_start = t->buf;
      t->buf_limit = t->buf;
      r = recv(t->sock_fd, t->buf_start, TESTIMONY_BUF_SIZE, 0);
      if (r == 0) {
        errno = ECANCELED;
        return -1;
      } else if (r < 0) {
        if (errno == EINTR) { continue; }
        return -1;
      }
      t->buf_limit = t->buf_start + r;
    }
    to_copy_out = out_limit - out;
    to_copy_buf = t->buf_limit - t->buf_start;
    to_copy = to_copy_out < to_copy_buf ? to_copy_out : to_copy_buf;
    memcpy(out, t->buf_start, to_copy);
    out += to_copy;
    t->buf_start += to_copy;
  }
  return 0;
}
static int recv_be_32(testimony t, uint32_t* out) {
  uint8_t msg[4];
  int r, i;
  r = recv_t(t, msg, 4);
  if (r < 0) { return r; }
  *out = 0;
  for (i = 0; i < 4; i++) {
    *out = (*out << 8) + msg[i];
  }
  return 0;
}

static int send_be_32(int fd, size_t in) {
  uint8_t msg[4];
  uint8_t* readfrom = msg;
  uint8_t* limit = msg + 4;
  int r;
  set_be_32(msg, in);
  while (readfrom < limit) {
    r = send(fd, readfrom, limit - readfrom, 0);
    if (r < 0) {
      return -1;
    }
    readfrom += r;
  }
  return 0;
}

// With much thanks to
// http://blog.varunajayasiri.com/passing-file-descriptors-between-processes-using-sendmsg-and-recvmsg
static int recv_file_descriptor(int socket) {
  struct msghdr message;
  struct iovec iov[1];
  struct cmsghdr* control_message = NULL;
  uint8_t ctrl_buf[CMSG_SPACE(sizeof(int))];
  uint8_t data[1];
  int r;

  memset(&message, 0, sizeof(struct msghdr));
  memset(ctrl_buf, 0, CMSG_SPACE(sizeof(int)));

  /* For the block data */
  iov[0].iov_base = data;
  iov[0].iov_len = sizeof(data);

  message.msg_name = NULL;
  message.msg_namelen = 0;
  message.msg_control = ctrl_buf;
  message.msg_controllen = CMSG_SPACE(sizeof(int));
  message.msg_iov = iov;
  message.msg_iovlen = 1;

  r = recvmsg(socket, &message, 0);
  if (r != sizeof(data)) {
    return -1;
  }

  /* Iterate through header to find if there is a file descriptor */
  for (control_message = CMSG_FIRSTHDR(&message); control_message != NULL;
       control_message = CMSG_NXTHDR(&message, control_message)) {
    if ((control_message->cmsg_level == SOL_SOCKET) &&
        (control_message->cmsg_type == SCM_RIGHTS)) {
      return *((int*)CMSG_DATA(control_message));
    }
  }

  errno = EIO;
  return -1;
}

static int discard_bytes(testimony t, size_t n) {
  uint8_t val[TESTIMONY_BUF_SIZE];
  size_t len;
  int r;
  while (n) {
    len = n > TESTIMONY_BUF_SIZE ? TESTIMONY_BUF_SIZE : n;
    r = recv_t(t, val, len);
    if (r < 0) {
      TERR("recv of val length %d", (int)n);
      return r;
    }
    n -= len;
  }
  return 0;
}

int testimony_connect(testimony* tp, const char* socket_name) {
  struct sockaddr_un saddr;
  sa_family_t laddr = AF_UNIX;  // Use unnamed socket on client side
  int r, err;
  uint8_t version;
  uint32_t msg, proto_typ, proto_len;
  testimony t = (testimony)malloc(sizeof(struct testimony_internal));
  if (t == NULL) {
    return -ENOMEM;
  }
  memset(t, 0, sizeof(*t));

  saddr.sun_family = AF_UNIX;
  strncpy(saddr.sun_path, socket_name, sizeof(saddr.sun_path) - 1);
  saddr.sun_path[sizeof(saddr.sun_path) - 1] = 0;

  t->sock_fd = socket(AF_UNIX, SOCK_STREAM, 0);
  if (t->sock_fd < 0) {
    TERR("socket creation failed");
    goto fail;
  }
  r = bind(t->sock_fd, (struct sockaddr*)&laddr, sizeof(laddr));
  if (r < 0) {
    TERR("bind failed");
    goto fail;
  }

  r = connect(t->sock_fd, (struct sockaddr*)&saddr, sizeof(saddr));
  if (r < 0) {
    TERR("connect to '%s' failed", saddr.sun_path);
    goto fail;
  }
  r = recv_t(t, &version, 1);
  if (r < 0) {
    TERR("recv of protocol version failed");
    goto fail;
  } else if (version != TESTIMONY_VERSION) {
    TERR_SET(EPROTONOSUPPORT, "received unsupported protocol version %d", version);
    goto fail;
  }
  proto_typ = TESTIMONY_PROTOCOL_TYPE_Error;
  while (1) {
    r = recv_be_32(t, &msg);
    if (r < 0) {
      TERR("did not receive fanout size");
      goto fail;
    }
    proto_typ = protocol_type(msg);
    proto_len = protocol_length(msg);
    if (proto_typ == TESTIMONY_PROTOCOL_TYPE_WaitingForFanoutIndex && proto_len == 0) {
      break;
    }
    if (proto_len == 4) {
      r = recv_be_32(t, &msg);
      if (r < 0) {
        TERR("did not receive type %d uint32 value", proto_typ);
        goto fail;
      }
      switch (proto_typ) {
        case TESTIMONY_PROTOCOL_TYPE_FanoutSize:
          t->conn.fanout_size = msg;
          break;
        case TESTIMONY_PROTOCOL_TYPE_BlockSize:
          t->conn.block_size = msg;
          break;
        case TESTIMONY_PROTOCOL_TYPE_NumBlocks:
          t->conn.block_nr = msg;
          break;
        default:
          // ignore
          break;
      }
    } else {
      r = discard_bytes(t, proto_len);
      if (r < 0) {
        TERR("did not discard type %d value", proto_typ);
        goto fail;
      }
    }
  }
  if (t->conn.fanout_size == 0 || t->conn.block_size == 0 || t->conn.block_nr == 0) {
    TERR_SET(EINVAL, "didn't get fanout size and block size/nr");
    goto fail;
  }
  *tp = t;
  return 0;
fail:
  err = errno;
  testimony_close(t);
  return -err;
}

testimony_connection* testimony_conn(testimony t) { return &t->conn; }

int testimony_init(testimony t) {
  uint8_t msg[4];
  int r;
  if (t->ring) {
    TERR_SET(EINVAL, "testimony has already been initiated");
    return -errno;
  }
  set_be_32(msg, (TESTIMONY_PROTOCOL_TYPE_FanoutIndex << 16) + 4);
  r = send(t->sock_fd, &msg, sizeof(msg), 0);
  if (r < 0) {
    TERR("send of fanout index type");
    return -errno;
  }
  set_be_32(msg, t->conn.fanout_index);
  r = send(t->sock_fd, &msg, sizeof(msg), 0);
  if (r < 0) {
    TERR("send of fanout index failed");
    return -errno;
  }

  t->afpacket_fd = recv_file_descriptor(t->sock_fd);
  if (t->afpacket_fd < 0) {
    TERR("recv of file descriptor failed");
    return -errno;
  }

  // calloc inits memory to zero
  t->block_counts = (uint32_t*)calloc(t->conn.block_nr, sizeof(uint32_t));

  t->ring = mmap(NULL, t->conn.block_size * t->conn.block_nr, PROT_READ,
                 MAP_SHARED | MAP_NORESERVE, t->afpacket_fd, 0);
  if (t->ring == MAP_FAILED) {
    t->ring = 0;
    TERR("local mmap of file descriptor failed");
    return -errno;
  }
  return 0;
}

int testimony_close(testimony t) {
  if (t->ring != 0) {
    if (munmap(t->ring, t->conn.block_nr * t->conn.block_size) < 0)
      return -errno;
  }
  if (close(t->sock_fd) < 0) return -errno;
  free(t->block_counts);
  free(t);
  return 0;
}

int testimony_get_block(testimony t, int timeout_millis,
                        const struct tpacket_block_desc** block) {
  // TODO:  Currently, testimony_get_block may block for longer than
  // timeout_millis, if either:
  //   1) the poll() is interrupted with a signal and restarts (EINTR handling)
  //   2) the recv (in recv_t) is unblocked by poll, but blocks before it can
  //      read in all the bytes it needs
  //   3) we block on the poll, read in a TLV, then block again
  // In short, we're really bad at using timeout_millis... should probably do
  // something about that.
  struct pollfd pfd;
  uint32_t blockidx, old_count;
  int r;
  uint32_t typ;
  *block = NULL;
  if (t->sock_fd == 0 || t->ring == 0) {
    TERR_SET(EINVAL, "testimony is not yet initiated, run testimony_init");
    return -errno;
  }
  while (1) {
    if (timeout_millis >= 0 && t->buf_start == t->buf_limit) {
      memset(&pfd, 0, sizeof(pfd));
      pfd.fd = t->sock_fd;
      pfd.events = POLLIN;
      r = poll(&pfd, 1, timeout_millis);
      if (r < 0) {
        if (errno == EINTR) { continue; }
        TERR("testimony poll of socket failed");
        return -errno;
      }
      if (r == 0) {
        return 0;  // Timed out, no block ready yet.
      }
      // A read is ready, fall through.
    }
    r = recv_be_32(t, &blockidx);
    if (r < 0) {
      TERR("recv of block index failed");
      return -errno;
    }
    typ = protocol_type(blockidx);
    if (typ == TESTIMONY_PROTOCOL_TYPE_BlockIndex) {
      break;
    }
    r = discard_bytes(t, protocol_length(blockidx));
    if (r < 0) {
      TERR("recv of block index failed");
      return -errno;
    }
  }
  if (blockidx >= t->conn.block_nr) {
    TERR_SET(EIO, "received invalid block index %d, should be [0, %d)",
             (int)blockidx, (int)t->conn.block_nr);
    return -errno;
  }
  *block =
      (const struct tpacket_block_desc*)(t->ring + t->conn.block_size * blockidx);
#ifdef __GCC_HAVE_SYNC_COMPARE_AND_SWAP_4
  if ((old_count = __sync_val_compare_and_swap(
      t->block_counts + blockidx, 0, (*block)->hdr.bh1.num_pkts)) != 0) {
    TERR_SET(EIO, "block count CAS failed for block %d, current count %d != 0",
             (int)blockidx, (int)old_count);
    return -errno;
  }
#endif
  return 0;
}

static const uint32_t kInvalidBlockIndex = 0xFFFFFFFF;

static uint32_t testimony_block_index(testimony t, const struct tpacket_block_desc* block) {
  size_t blockptr = (size_t)block;
  blockptr -= (size_t)t->ring;
  blockptr /= t->conn.block_size;
  if (blockptr >= t->conn.block_nr) {
    return kInvalidBlockIndex;
  }
  return blockptr;
}

int testimony_return_block(testimony t, const struct tpacket_block_desc* block) {
  int r;
  uint32_t old_count;
  uint32_t blockidx = testimony_block_index(t, block);
  if (blockidx == kInvalidBlockIndex) {
    TERR_SET(EINVAL, "block does not appear to have come from this testimony instance");
    return -errno;
  }
#ifdef __GCC_HAVE_SYNC_COMPARE_AND_SWAP_4
  // Set block count for this block to zero (& with 0), and make sure the packet
  // count was sane.
  old_count = __sync_fetch_and_and(t->block_counts + blockidx, 0);
  if (old_count != 0 && old_count != block->hdr.bh1.num_pkts) {
    TERR_SET(EINVAL, "block count invalid... maybe testimony_return_block and "
             "testimony_return_packet were both called?");
    return -errno;
  }
#endif
  r = send_be_32(t->sock_fd, blockidx);
  if (r < 0) {
    TERR("send of block index failed");
    return -errno;
  }
  return 0;
}

int testimony_return_packets(testimony t, const struct tpacket_block_desc* block, uint32_t packets) {
#ifdef __GCC_HAVE_SYNC_COMPARE_AND_SWAP_4
  uint32_t blockidx = testimony_block_index(t, block);
  uint32_t count;
  if (blockidx == kInvalidBlockIndex) {
    TERR_SET(EINVAL, "block does not appear to have come from this testimony instance");
    return -errno;
  }
  count = __sync_fetch_and_sub(t->block_counts + blockidx, packets);
  if (count == packets) {
    return testimony_return_block(t, block);
  } else if (count < packets) {
    TERR_SET(EINVAL, "return_packets amount overflows number of packets");
    return -errno;
  }
  return 0;
#else
  TERR_SET(EINVAL, "compiler does not support atomics, "
           "testimony_return_packet not available");
  return -errno;
#endif
}

char* testimony_error(testimony t) { return t->errbuf; }

struct testimony_iter_internal {
  const struct tpacket_block_desc* block;
  uint8_t* pkt;
  int left;
};

int testimony_iter_init(testimony_iter* iter) {
  *iter = (testimony_iter)malloc(sizeof(struct testimony_iter_internal));
  if (*iter == NULL) {
    return -ENOMEM;
  }
  memset(*iter, 0, sizeof(struct testimony_iter_internal));
  return 0;
}

int testimony_iter_reset(testimony_iter iter,
                         const struct tpacket_block_desc* block) {
  if (block->version != TPACKET_V3) {
    return -EPROTONOSUPPORT;
  }
  iter->block = block;
  iter->left = block->hdr.bh1.num_pkts;
  iter->pkt = NULL;
  return 0;
}

const struct tpacket3_hdr* testimony_iter_next(testimony_iter iter) {
  if (!iter->left) {
    return NULL;
  }
  iter->left--;
  if (iter->pkt) {
    iter->pkt += ((const struct tpacket3_hdr*)iter->pkt)->tp_next_offset;
  } else {
    iter->pkt =
        (uint8_t*)iter->block + iter->block->hdr.bh1.offset_to_first_pkt;
  }
  return (const struct tpacket3_hdr*)iter->pkt;
}

int testimony_iter_close(testimony_iter iter) {
  free(iter);
  return 0;
}

const uint8_t* testimony_packet_data(const struct tpacket3_hdr* pkt) {
  return (const uint8_t*)pkt + pkt->tp_mac;
}
int64_t testimony_packet_nanos(const struct tpacket3_hdr* pkt) {
  return (int64_t)pkt->tp_sec * 1000000000LL + pkt->tp_nsec;
}

#ifdef __cplusplus
}
#endif
